/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;


import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

import alloy.ast.ASTDepthFirstReplacer;
import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.ArrowMultExpr;
import alloy.ast.CompOp;
import alloy.ast.Decl;
import alloy.ast.Decls;
import alloy.ast.ElemFormula;
import alloy.ast.Expr;
import alloy.ast.Formula;
import alloy.ast.FormulaSeq;
import alloy.ast.Formulas;
import alloy.ast.Id;
import alloy.ast.Multiplicity;
import alloy.ast.MultiplicityExpr;
import alloy.ast.MultiplicityFormula;
import alloy.ast.Qualifiers;
import alloy.ast.QuantifiedExpr;
import alloy.ast.QuantifiedFormula;
import alloy.ast.Quantifier;
import alloy.ast.SetMultExpr;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.ast.Variables;
import alloy.type.RelationType;
import alloy.type.UnionType;

/**
 * <p>A visitor which desugars multiplicity formulas.  A set multiplicity formula is 
 * desugared as follows:
 * <ul>
 * <li> x : set expr or x : expr where arity(expr)>1 ==> x : expr
 * <li> x : (lone | one | some) expr ==> x : expr && (lone | one | some) x
 * <li> x : expr where arity(expr)=1 ==> x : expr && one x
 * </ul>
 * <p>An arrow multiplicity formula has the following nice desugaring due to gdennis:
 * I'll have to define a function L(v) that returns the sequence of leaf
 * expressions in the multiplicity expression v in order. that is, if we
 * have a multiplicity expression like:
 *
 *  r: (E1 m->n E2) p->q E3
 *
 * then
 *
 *  L((E1 m->n E2) p->q E3) = E1, E2, E3
 *  L(E1 m->n E2) = E1, E2
 *  L(E3) = E3
 *
 * in the formulation of the desugaring below, I've assumed have the
 * ability to use all set multiplicities (no, one, lone, some) on relations
 * to constrain the number of tuples in the relation. that is, assume we
 * can write formulas like "lone A->B" to say there's at most one tuple in
 * the cross product of A and B. if we don't, we can rewrite them in terms
 * of syntax we support.
 * 
 * Assume that type[Xi] = Xi_1->Xi_2->...->Xi_arity(Xi) and type[Yi] = Yi_1->Yi_2->...->Yi_arity(Yi)
 *
 *
 *  T[r: m X p->q Y], L(X) = X1,...,Xn, L(Y) = Y1,...,Yk
 * ======================================================>
 *  m r &&
 *  all x1_1: X1_1,...,x1_arity(X1): X1_ai, ..., xn_1 : Xn_1, ..., xn_arity(Xn) : Xn_arity(Xn) | 
 *    (x1_1->...->x1_arity(X1) in X1 && ... && xn_1->...->xn_arity(Xn) in Xn) => 
 *    T[r[x1_1] ... [xn_arity(Xn)] : q Y]
 *  all y1_1: Y1_1,...,y1_arity(Y1): Y1_bi, ..., yk_1 : Yk_1, ..., yk_arity(Yk) : Yk_arity(Yk) | 
 *    (y1_1->...->y1_arity(Y1) in Y1 && ... && yk_1->...->yk_arity(Yk) in Yk) => 
 *    T[r.yk_arity(Yk). ... .y1_1  : p X]
 *
 */
public class DesugarMultFormulasVisitor extends ASTDepthFirstReplacer {
    private static final String VAR_NAME_BASE = "var";
    
    /** 
     * Expression that is constrained in the MultiplicityFormula we are 
     * currently desugaring.  This field is null if we are not inside
     * a MultiplicityFormula at the moment.
     **/
    private Expr curExpr = null;

	/**
 	 * @param mult_ determines multiplicity and location
 	 * @param expr_ Expr expression
 	 * @return the {@link Formula} "(lone | one | some) expr_". The multiplicity of
	 * the generated formula is determined by mult_. If mult_ is
	 * Multiplicity.NONE and the type of expr_ has arity 1, then NONE is treated
	 * as ONE; otherwise, it's treated as SET. Null is returned if no
	 * multiplicity constraint applies to the expression. The location of the
	 * returned formula is set to the location of mult_.
 	 */
    public static Formula genSetMultConstraint(Multiplicity mult_, Expr expr_) {
		// x : S means x : one S when S is a set, otherwise it means x : set S . . . 
		if (mult_.equals(Multiplicity.NONE)) {
		    mult_ = ( expr_.getType().arity()==1 ? Multiplicity.ONE : Multiplicity.SET);  
		}
	
		Formula quantExpr = null;
	
		if (!mult_.equals(Multiplicity.SET)) {
		    if      (mult_.equals(Multiplicity.LONE)) quantExpr = QuantifiedExpr.lone((Expr)expr_.copy());
		    else if (mult_.equals(Multiplicity.ONE))  quantExpr = QuantifiedExpr.one((Expr)expr_.copy());
		    else if (mult_.equals(Multiplicity.SOME)) quantExpr = QuantifiedExpr.some((Expr)expr_.copy());
		  	    
		    quantExpr.annotate(mult_.nodeString() + " multiplicity constraint");
		    quantExpr.setLocation(mult_.getLocation());
		}
	
		return quantExpr;
    }

    public Object visit(SetMultExpr setMultExpr_) {
	if (curExpr == null) return setMultExpr_; // nothing to do ...

	Formulas formulas = new Formulas();

        // constraint for set membership
        ElemFormula membershipFormula = new ElemFormula((Expr) curExpr.copy(),
							CompOp.SUBSETIN,
							(Expr) setMultExpr_.getExpr().copy());

        membershipFormula.annotate("set membership constraint");
        membershipFormula.setLocation(curExpr.getLocation());

        formulas.addFormula(membershipFormula);

		// constraint for multiplicity
		Formula multConstraint = genSetMultConstraint(setMultExpr_.getMult(), curExpr);
		if (multConstraint != null)
			formulas.addFormula(multConstraint);
	
        return new FormulaSeq(formulas);
    }

    /** 
     * Assumes that exprList is a list of expressions, [X1, X2, ..., Xn]
     * If leftAssociate, we generate the formula
     *  all v1_1: X1_1, ... , v1_arity(X1): X1_arity(X1), ... vn_1: Xn_1, ..., vn_arity(Xn): Xn_arity(Xn) | 
     *    (v1_1->...->v1_arity(X1) in X1 && ... && vn_1->...->vn_arity(Xn) in Xn) => 
     *    T[(...((expr.v1_1).v1_2).v1_3 ... ).vn_arity(Xn) : multExpr]
     * If !leftAssociate, we generate the formula
     *  all v1_1: X1_1, ... , v1_arity(X1): X1_arity(X1), ... vn_1: Xn_1, ..., vn_arity(Xn): Xn_arity(Xn) | 
     *    (v1_1->...->v1_arity(X1) in X1 && ... && vn_1->...->vn_arity(Xn) in Xn) => 
     *    T[v1_1.(... .(vn_arity(Xn).expr))...) : multExpr]
     * Above formulas assume that type[Xi] = Xi_1->Xi_2->...->Xi_arity(Xi)
     **/ 
    private static Formula arrowMultConstraint(List exprList, MultiplicityExpr multExpr, 
					       Expr expr, boolean leftAssociate) {

	Decls decls = new Decls();
	Expr joinedExpr = (Expr) expr.copy();
	Formula antecedent = null;

	for (Iterator exprIter = exprList.iterator(); exprIter.hasNext();) {
	    Expr declExpr = (Expr) ((Expr) exprIter.next()).copy();
	    
	    Expr productExpr = null;

	    // vi_1: Xi_1, ... ,vi_arity(Xi): Xi_arity(Xi)
	    UnionType declType = declExpr.getType();

	    
	    for (int i = 0; i < declType.arity(); i++) {
	    
		Variables vars = new Variables();
		Variable var = new Variable(Id.generateId(VAR_NAME_BASE));

		vars.addVariable(var);

		UnionType varType = UnionType.empty(1);
		for (Iterator iter = declType.getTypeIter(); iter.hasNext();) {
		    RelationType relType = (RelationType) iter.next();
		    varType = varType.union(new UnionType(new RelationType(relType.getBasicTypeAt(i))));
		}
	    	Decl decl = new Decl(declExpr.getLocation(), new Qualifiers(), vars, 
	    	        					new SetMultExpr(declExpr.getLocation(),	
	    	        		        		Multiplicity.NONE, varType.toExpr()));   
	    	decl.desugared = true;
		decls.addDecl(decl);

		VariableExpr varExpr = new VariableExpr((Variable) var.copy());
		varExpr.setType(varType);
		
		if (leftAssociate) // (...((expr.v1).v2).v3 ... ).vn
		    joinedExpr = joinedExpr.join(varExpr);
		else // v1.(... vn-2.(vn-1.(vn.expr))...)
		    joinedExpr = varExpr.join(joinedExpr);

		// vi_1->...->vi_arity(Xi)
		productExpr = (productExpr == null ? (Expr) varExpr.copy() : 
			       productExpr.product((Expr) varExpr.copy()));
	
	    }

	    // v1_1->...->v1_arity(X1) in X1
	    antecedent = (antecedent == null ? productExpr.in((Expr)declExpr.copy()) : 
			  antecedent.and(productExpr.in((Expr)declExpr.copy())));

	}
	
	// in an arrow multiplicity expr, arrow multiplicity of NONE means SET
	// e.g. X -> Y means X set->set Y
	MultiplicityExpr multExprCopy = (MultiplicityExpr) multExpr.copy();
	if (multExprCopy.getMult().equals(Multiplicity.NONE))
	    multExprCopy.setMult(Multiplicity.SET);
	

	// T[joinedExpr : multExpr]
	Formula multFormula = new MultiplicityFormula(multExpr.getLocation(), joinedExpr, 
						      multExprCopy);

	multFormula = (Formula) multFormula.applyReturnVisitor(new DesugarMultFormulasVisitor());
		
	//  all v1_1: X1_1, ... , v1_arity(X1): X1_arity(X1), ... vn_1: Xn_1, ..., vn_arity(Xn): Xn_arity(Xn) | 
	//    (v1_1->...->v1_arity(X1) in X1 && ... && vn_1->...->vn_arity(Xn) in Xn) => 
	//    T[joinedExpr : multExpr]

	return new QuantifiedFormula(
			Quantifier.ALL,
			decls,
			(antecedent == null ? multFormula : antecedent.implies(multFormula)),
			true);

    }

    /**
     *  T[r: m X p->q Y], L(X) = X1,...,Xn, L(Y) = Y1,...,Yk
     * ======================================================>
     *  m r &&
     *  all x1: X1,...,xn: Xn | T[r[x1] ... [xn] : q Y]
     *  all y1: Y1,...,yk: Yk | T[r.yk. ... .y1  : p X]
     **/
    public Object visit(ArrowMultExpr arrowMultExpr_) {
	
	if (curExpr == null) return arrowMultExpr_; // nothing to do
	
	Formulas formulas = new Formulas();

	// m r
	Formula multConstraint = genSetMultConstraint(arrowMultExpr_.getMult(), curExpr);
	if (multConstraint != null) formulas.addFormula(multConstraint);

	// r : X->Y
	formulas.addFormula(((Expr)curExpr.copy()).in(arrowMultExpr_.toExpr()));

	MultiplicityExpr left = arrowMultExpr_.getLeft();
	MultiplicityExpr right = arrowMultExpr_.getRight();


	// all x1: X1,...,xn: Xn | T[xn.(...(x2.(x1.r))...) : q Y]
	formulas.addFormula(arrowMultConstraint(ExprCollector.exprSeq(left), 
						right, curExpr, false)); 
							   
	
	// all y1: Y1,...,yk: Yk | T[r.yk. ... .y1  : p X]
	formulas.addFormula(arrowMultConstraint(ExprCollector.reversedExprSeq(right), 
						left, curExpr, true));
	
	return new FormulaSeq(formulas);
    }

	public Object visit(MultiplicityFormula multFormula_) {
		curExpr = multFormula_.getLeft();
		Formula formula = (Formula) multFormula_.getRight().applyReturnVisitor(this);
		formula.setLocation(multFormula_.getLocation());
		curExpr = null;

		return formula;
	}

    /**
     * Computes the L(X) function defined at the beginning of the file.
     * i.e.  L((E1 m->n E2) p->q E3) = E1, E2, E3
     **/
	private static class ExprCollector extends ASTDepthFirstVisitor {
		private List exprs = new LinkedList();

		private ExprCollector() {
		}

		static List exprSeq(MultiplicityExpr multExpr_) {
			ExprCollector collector = new ExprCollector();
			multExpr_.applyVisitor(collector);
			return collector.exprs;
		}

		static List reversedExprSeq(MultiplicityExpr multExpr_) {
			List seq = exprSeq(multExpr_);
			Collections.reverse(seq);
			return seq;
		}

		public void visit(SetMultExpr setMultExpr_) {
			exprs.add(setMultExpr_.getExpr());
		}
	}
	
}
